import asyncio
import time
import random
from typing import Dict, Any
from loguru import logger

from app.core.resource_scheduler import (
    get_resource_scheduler, 
    TaskPriority,
    RetryConfig
)


# 配置日志
logger.add("resource_scheduler_example.log", rotation="10 MB")


async def example_task(task_name: str, task_params: Dict[str, Any]) -> Dict[str, Any]:
    """
    示例任务函数
    
    Args:
        task_name: 任务名称
        task_params: 任务参数
        
    Returns:
        任务结果
    """
    logger.info(f"执行任务: {task_name}, 参数: {task_params}")
    
    # 模拟任务执行时间
    execution_time = task_params.get('execution_time', random.uniform(0.5, 3.0))
    await asyncio.sleep(execution_time)
    
    # 模拟一些资源消耗
    if task_params.get('simulate_heavy_load', False):
        # 模拟CPU密集型操作
        for _ in range(1000000):
            _ = random.random() * random.random()
    
    # 模拟任务失败
    if task_params.get('should_fail', False):
        raise ValueError(f"模拟任务 {task_name} 失败")
    
    result = {
        'task_name': task_name,
        'status': 'success',
        'execution_time': execution_time,
        'timestamp': time.time(),
        'params': task_params
    }
    
    logger.info(f"任务 {task_name} 执行完成，耗时: {execution_time:.2f}秒")
    return result


def sync_example_task(task_name: str, task_params: Dict[str, Any]) -> Dict[str, Any]:
    """
    同步示例任务函数
    
    Args:
        task_name: 任务名称
        task_params: 任务参数
        
    Returns:
        任务结果
    """
    logger.info(f"执行同步任务: {task_name}, 参数: {task_params}")
    
    # 模拟任务执行时间
    execution_time = task_params.get('execution_time', random.uniform(0.5, 2.0))
    time.sleep(execution_time)
    
    # 模拟一些资源消耗
    if task_params.get('simulate_heavy_load', False):
        for _ in range(500000):
            _ = random.random() * random.random()
    
    # 模拟任务失败
    if task_params.get('should_fail', False):
        raise ValueError(f"模拟同步任务 {task_name} 失败")
    
    result = {
        'task_name': task_name,
        'status': 'success',
        'execution_time': execution_time,
        'timestamp': time.time(),
        'params': task_params
    }
    
    logger.info(f"同步任务 {task_name} 执行完成，耗时: {execution_time:.2f}秒")
    return result


async def monitor_scheduler(scheduler):
    """
    监控调度器状态
    
    Args:
        scheduler: 资源调度器实例
    """
    while True:
        queue_status = scheduler.get_queue_status()
        stats = scheduler.get_statistics()
        
        logger.info(f"调度器状态: 运行中任务={queue_status['running_tasks']}, "
                   f"队列任务={queue_status['queue_length']}, "
                   f"成功率={stats['success_rate']:.1f}%")
        
        await asyncio.sleep(5)  # 每5秒监控一次


async def run_example():
    """
    运行资源调度器示例
    """
    # 获取资源调度器实例
    scheduler = get_resource_scheduler()
    
    # 启动监控任务
    monitor_task = asyncio.create_task(monitor_scheduler(scheduler))
    
    try:
        # 创建重试配置
        retry_config = RetryConfig(
            max_retries=3,
            base_delay=1.0,
            max_delay=5.0,
            backoff_factor=2.0,
            jitter=True
        )
        
        # 调度不同优先级的任务
        logger.info("开始调度任务...")
        
        # 1. 关键任务 - 立即执行
        task_ids = []
        task_id = await scheduler.schedule_task(
            task_id="critical_task_001",
            task_func=example_task,
            priority=TaskPriority.CRITICAL,
            metadata={
                "task_name": "系统健康检查",
                "execution_time": 1.0
            },
            resource_requirements={
                "cpu_cores": 1,
                "memory_mb": 100
            }
        )
        task_ids.append(task_id)
        
        # 2. 高优先级任务
        for i in range(3):
            task_id = await scheduler.schedule_task(
                task_id=f"high_task_{i:03d}",
                task_func=example_task,
                priority=TaskPriority.HIGH,
                metadata={
                    "task_name": f"用户请求处理_{i}",
                    "execution_time": random.uniform(0.5, 2.0),
                    "user_id": f"user_{i}"
                },
                resource_requirements={
                    "cpu_cores": 2,
                    "memory_mb": 200
                }
            )
            task_ids.append(task_id)
        
        # 3. 中优先级任务 - 包含一个会失败并重试的任务
        task_id = await scheduler.schedule_task(
            task_id="medium_task_fail_retry",
            task_func=example_task,
            priority=TaskPriority.MEDIUM,
            metadata={
                "task_name": "数据分析任务",
                "execution_time": 1.5,
                "should_fail": True  # 第一次执行会失败
            },
            resource_requirements={
                "cpu_cores": 2,
                "memory_mb": 300
            },
            retry_config=retry_config
        )
        task_ids.append(task_id)
        
        # 4. 低优先级任务 - 模拟重负载任务
        for i in range(2):
            task_id = await scheduler.schedule_task(
                task_id=f"low_task_{i:03d}",
                task_func=sync_example_task,  # 使用同步函数
                priority=TaskPriority.LOW,
                metadata={
                    "task_name": f"数据备份_{i}",
                    "execution_time": random.uniform(2.0, 4.0),
                    "simulate_heavy_load": True
                },
                resource_requirements={
                    "cpu_cores": 4,
                    "memory_mb": 500
                }
            )
            task_ids.append(task_id)
        
        # 5. 后台任务
        task_id = await scheduler.schedule_task(
            task_id="background_task_001",
            task_func=example_task,
            priority=TaskPriority.BACKGROUND,
            metadata={
                "task_name": "日志清理",
                "execution_time": 3.0
            },
            resource_requirements={
                "cpu_cores": 1,
                "memory_mb": 150
            }
        )
        task_ids.append(task_id)
        
        # 测试取消任务
        cancel_task_id = await scheduler.schedule_task(
            task_id="cancel_test_task",
            task_func=example_task,
            priority=TaskPriority.LOW,
            metadata={
                "task_name": "测试取消任务",
                "execution_time": 10.0  # 长时间运行的任务
            }
        )
        
        # 等待一小段时间后取消任务
        await asyncio.sleep(1)
        cancelled = await scheduler.cancel_task(cancel_task_id)
        logger.info(f"任务取消结果: {cancelled}, 任务ID: {cancel_task_id}")
        
        # 等待所有任务完成
        logger.info(f"已调度 {len(task_ids)} 个任务，等待执行完成...")
        
        # 等待足够长的时间让所有任务完成
        await asyncio.sleep(30)
        
        # 获取并打印任务状态
        logger.info("\n任务执行结果:")
        for task_id in task_ids:
            status = scheduler.get_task_status(task_id)
            if status:
                logger.info(f"任务 {task_id}: 状态={status['status']}, "
                           f"优先级={status['priority']}, "
                           f"耗时={status.get('execution_time', 'N/A'):.2f}秒")
        
        # 打印调度器统计信息
        logger.info("\n调度器统计信息:")
        stats = scheduler.get_statistics()
        logger.info(f"总任务数: {stats['total_tasks']}")
        logger.info(f"完成任务数: {stats['completed_tasks']}")
        logger.info(f"失败任务数: {stats['failed_tasks']}")
        logger.info(f"平均执行时间: {stats['average_execution_time']:.2f}秒")
        logger.info(f"成功率: {stats['success_rate']:.1f}%")
        logger.info(f"平均资源使用率: {stats['average_resource_usage']:.1f}%")
        
    except Exception as e:
        logger.error(f"示例运行出错: {str(e)}")
    finally:
        # 取消监控任务
        monitor_task.cancel()
        try:
            await monitor_task
        except asyncio.CancelledError:
            pass
        
        # 关闭调度器
        await scheduler.stop()
        logger.info("资源调度器示例运行完成")


if __name__ == "__main__":
    try:
        asyncio.run(run_example())
    except KeyboardInterrupt:
        logger.info("程序被用户中断")
    except Exception as e:
        logger.error(f"程序运行出错: {str(e)}")
